import tkinter as tk
from tkinter import scrolledtext, messagebox, filedialog, ttk, simpledialog
import requests
import json
import os
from datetime import datetime
import markdown
from tkhtmlview import HTMLLabel

class ChatApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Chat App")
        self.root.geometry("1400x900")  # 增加窗口大小
        
        # 添加配置文件路径
        self.config_file = "chat_app_config.json"
        
        # 初始化配置属性
        self.saved_workspace = None
        self.saved_thread = None
        self.current_workspace = None
        
        # 添加默认工作区
        self.default_workspace = "default"
        self.api_base_url = "http://localhost:3001"  # 添加API基础URL

        # 创建主框架
        main_frame = tk.PanedWindow(root, orient=tk.HORIZONTAL)
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 左侧面板
        left_frame = tk.Frame(main_frame)
        main_frame.add(left_frame)

        # 右侧面板（聊天区域）
        right_frame = tk.Frame(main_frame)
        main_frame.add(right_frame)

        # 创建标签页
        self.notebook = ttk.Notebook(left_frame)
        self.notebook.pack(fill=tk.BOTH, expand=True)

        # 设置标签页
        settings_frame = ttk.Frame(self.notebook)
        doc_frame = ttk.Frame(self.notebook)
        self.notebook.add(settings_frame, text="设置")
        self.notebook.add(doc_frame, text="文档管理")

        # API设置区域
        api_frame = tk.LabelFrame(settings_frame, text="API 设置", padx=5, pady=5)
        api_frame.pack(fill=tk.X, padx=5, pady=5)

        tk.Label(api_frame, text="API 密钥:").pack(anchor="w")
        self.api_key_entry = tk.Entry(api_frame, show="*")
        self.api_key_entry.pack(fill=tk.X)
        self.api_key_entry.insert(0, "F5CMH7P-QYH4Q4F-Q5M3SPJ-FZJZKJ1")  # 默认API key

        tk.Label(api_frame, text="API 地址:").pack(anchor="w")
        self.api_url_entry = tk.Entry(api_frame)
        self.api_url_entry.pack(fill=tk.X)
        self.api_url_entry.insert(0, self.api_base_url)
        tk.Button(api_frame, text="更新地址", command=self.update_api_url).pack(fill=tk.X, pady=2)

        # 工作区管理
        workspace_frame = tk.LabelFrame(settings_frame, text="工作区管理", padx=5, pady=5)
        workspace_frame.pack(fill=tk.X, padx=5, pady=5)

        default_ws_frame = tk.Frame(workspace_frame)
        default_ws_frame.pack(fill=tk.X)
        tk.Label(default_ws_frame, text="默认工作区:").pack(side=tk.LEFT)
        self.default_ws_entry = tk.Entry(default_ws_frame)
        self.default_ws_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.default_ws_entry.insert(0, self.default_workspace)
        tk.Button(default_ws_frame, text="设为默认", command=self.set_default_workspace).pack(side=tk.LEFT)
        tk.Button(default_ws_frame, text="保存对话", command=self.save_chat_history).pack(side=tk.LEFT, padx=5)

        tk.Button(workspace_frame, text="创建工作区", command=self.create_workspace).pack(fill=tk.X)
        tk.Button(workspace_frame, text="刷新工作区", command=self.refresh_workspaces).pack(fill=tk.X)
        self.workspace_listbox = tk.Listbox(workspace_frame, height=5)
        self.workspace_listbox.pack(fill=tk.X)
        tk.Button(workspace_frame, text="删除工作区", command=self.delete_workspace).pack(fill=tk.X)

        # 文档管理区域
        doc_manage_frame = tk.LabelFrame(doc_frame, text="文档管理", padx=5, pady=5)
        doc_manage_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        upload_frame = tk.Frame(doc_manage_frame)
        upload_frame.pack(fill=tk.X)

        tk.Button(upload_frame, text="上传文件", command=self.select_file).pack(side=tk.LEFT)
        self.file_label = tk.Label(upload_frame, text="未选择文件")
        self.file_label.pack(side=tk.LEFT, padx=5)

        tk.Label(doc_manage_frame, text="网站地址:").pack(anchor="w")
        self.website_entry = tk.Entry(doc_manage_frame)
        self.website_entry.pack(fill=tk.X)
        tk.Button(doc_manage_frame, text="上传网站", command=self.upload_website).pack(fill=tk.X)

        self.doc_listbox = tk.Listbox(doc_manage_frame, height=10)
        self.doc_listbox.pack(fill=tk.BOTH, expand=True)
        
        # 文档管理按钮框架
        doc_buttons_frame = tk.Frame(doc_manage_frame)
        doc_buttons_frame.pack(fill=tk.X)
        
        tk.Button(doc_buttons_frame, text="刷新文档", command=self.refresh_documents).pack(side=tk.LEFT, fill=tk.X, expand=True)
        tk.Button(doc_buttons_frame, text="删除文档", command=self.delete_document).pack(side=tk.LEFT, fill=tk.X, expand=True)
        tk.Button(doc_buttons_frame, text="更新向量", command=self.update_embeddings).pack(side=tk.LEFT, fill=tk.X, expand=True)

        # 对话区域
        chat_frame = tk.LabelFrame(right_frame, text="对话", padx=5, pady=5)
        chat_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # Thread管理
        thread_frame = tk.Frame(chat_frame)
        thread_frame.pack(fill=tk.X)
        
        tk.Label(thread_frame, text="对话用户:").pack(side=tk.LEFT)
        self.thread_entry = tk.Entry(thread_frame)
        self.thread_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.thread_entry.insert(0, "default")  # 默认线程名
        tk.Button(thread_frame, text="创建用户", command=self.create_thread).pack(side=tk.LEFT)
        tk.Button(thread_frame, text="加载历史", command=self.load_chat_history).pack(side=tk.LEFT, padx=5)

        # 添加模式选择
        mode_frame = tk.Frame(chat_frame)
        mode_frame.pack(fill=tk.X)
        tk.Label(mode_frame, text="模式:").pack(side=tk.LEFT)
        self.mode_var = tk.StringVar(value="query")
        tk.Radiobutton(mode_frame, text="问答", variable=self.mode_var, value="query").pack(side=tk.LEFT)
        tk.Radiobutton(mode_frame, text="聊天", variable=self.mode_var, value="chat").pack(side=tk.LEFT)

        # 聊天显示区域
        self.chat_display = scrolledtext.ScrolledText(chat_frame, height=30)
        self.chat_display.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 配置文本标签
        self.chat_display.tag_configure("user", foreground="blue")
        
        # 消息输入区域
        input_frame = tk.Frame(chat_frame)
        input_frame.pack(fill=tk.X, padx=5, pady=5)
        
        self.question_entry = tk.Text(input_frame, height=3)
        self.question_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 绑定键盘事件
        self.question_entry.bind('<Return>', self.on_enter)
        self.question_entry.bind('<Shift-Return>', self.on_shift_enter)
        
        # 发送按钮
        send_button = tk.Button(input_frame, text="发送", command=self.send_message, width=10)
        send_button.pack(side=tk.LEFT, padx=5)

        # 加载保存的配置
        self.load_config()
        
        # 在窗口关闭时保存配置
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        # 初始化数据
        self.root.after(1000, self.initialize_workspace)  # 延迟1秒初始化

    def save_config(self):
        """保存配置到文件"""
        config = {
            "workspace": self.current_workspace,
            "thread": self.thread_entry.get(),
            "api_key": self.api_key_entry.get(),
            "api_url": self.api_url_entry.get()
        }
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"Error saving config: {str(e)}")

    def load_config(self):
        """从文件加载配置"""
        # 设置默认值
        self.saved_workspace = self.default_workspace
        self.saved_thread = "default"
        
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    
                    # 恢复API设置
                    if "api_key" in config:
                        self.api_key_entry.delete(0, tk.END)
                        self.api_key_entry.insert(0, config["api_key"])
                    if "api_url" in config:
                        self.api_url_entry.delete(0, tk.END)
                        self.api_url_entry.insert(0, config["api_url"])
                        self.api_base_url = config["api_url"]
                    
                    # 加载工作区和线程信息
                    self.saved_workspace = config.get("workspace", self.default_workspace)
                    self.saved_thread = config.get("thread", "default")
                    
            # 如果配置文件不存在，使用默认值
            if not os.path.exists(self.config_file):
                self.save_config()  # 创建默认配置文件
                
        except Exception as e:
            print(f"Error loading config: {str(e)}")
            # 发生错误时使用默认值
            self.saved_workspace = self.default_workspace
            self.saved_thread = "default"
            
        # 设置默认线程名
        self.thread_entry.delete(0, tk.END)
        self.thread_entry.insert(0, self.saved_thread)

    def on_closing(self):
        """窗口关闭时的处理"""
        self.save_config()
        self.root.destroy()

    def initialize_workspace(self):
        """初始化工作区"""
        try:
            # 首先检查默认工作区是否存在
            api_key = self.api_key_entry.get()
            url = f"{self.api_base_url}/api/v1/workspaces"
            headers = {"Authorization": f"Bearer {api_key}"}
            
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                data = response.json()
                workspaces = data.get("workspaces", [])
                workspace_names = [ws["name"] for ws in workspaces]
                
                # 如果默认工作区不存在，创建它
                if self.default_workspace not in workspace_names:
                    self.create_default_workspace()
                else:
                    self.current_workspace = self.default_workspace
                    self.refresh_workspaces()
            else:
                messagebox.showerror("Error", f"Failed to fetch workspaces: {response.text}")
        except Exception as e:
            messagebox.showwarning("Warning", f"Failed to initialize workspace: {str(e)}")

    def create_workspace(self):
        dialog = WorkspaceConfigDialog(self.root)
        self.root.wait_window(dialog)
        
        if dialog.result:
            api_key = self.api_key_entry.get()
            url = f"{self.api_base_url}/api/v1/workspace/new"
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            try:
                response = requests.post(url, headers=headers, json=dialog.result)
                if response.status_code == 200:
                    messagebox.showinfo("Success", "Workspace created successfully!")
                    self.refresh_workspaces()
                else:
                    messagebox.showerror("Error", f"Failed to create workspace: {response.text}")
            except Exception as e:
                messagebox.showerror("Error", str(e))

    def create_default_workspace(self):
        """创建默认工作区"""
        api_key = self.api_key_entry.get()
        url = f"{self.api_base_url}/api/v1/workspace/new"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        data = {
            "name": self.default_workspace,
            "similarityThreshold": 0.7,
            "openAiTemp": 0.7,
            "openAiHistory": 20,
            "openAiPrompt": "Given the following conversation, relevant context, and a follow up question, reply with an answer to the current question the user is asking. Return only your response to the question given the above information following the users instructions as needed.",
            "queryRefusalResponse": "There is no relevant information in this workspace to answer your query.",
            "chatMode": "chat",
            "topN": 4
        }
        
        try:
            response = requests.post(url, headers=headers, json=data)
            if response.status_code == 200:
                self.current_workspace = self.default_workspace
                self.refresh_workspaces()
                self.refresh_documents()
            else:
                messagebox.showerror("Error", f"Failed to create default workspace: {response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def set_default_workspace(self):
        """设置默认工作区"""
        new_default = self.default_ws_entry.get()
        if new_default:
            self.default_workspace = new_default
            self.initialize_workspace()

    def refresh_workspaces(self):
        api_key = self.api_key_entry.get()
        url = f"{self.api_base_url}/api/v1/workspaces"
        headers = {"Authorization": f"Bearer {api_key}"}
        
        try:
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                data = response.json()
                workspaces = data.get("workspaces", [])
                self.workspace_listbox.delete(0, tk.END)
                
                # 填充工作区列表
                for workspace in workspaces:
                    self.workspace_listbox.insert(tk.END, workspace["name"])
                
                # 如果有保存的工作区，选择它
                if self.saved_workspace and self.workspace_listbox.size() > 0:
                    for i, workspace in enumerate(workspaces):
                        if workspace["name"] == self.saved_workspace:
                            self.workspace_listbox.selection_clear(0, tk.END)
                            self.workspace_listbox.selection_set(i)
                            self.current_workspace = self.saved_workspace
                            # 恢复保存的线程
                            if self.saved_thread:
                                self.thread_entry.delete(0, tk.END)
                                self.thread_entry.insert(0, self.saved_thread)
                            break
                    else:  # 如果找不到保存的工作区，选择第一个
                        self.workspace_listbox.selection_clear(0, tk.END)
                        self.workspace_listbox.selection_set(0)
                        self.current_workspace = self.workspace_listbox.get(0)
                elif self.workspace_listbox.size() > 0:  # 如果没有保存的工作区，选择第一个
                    self.workspace_listbox.selection_clear(0, tk.END)
                    self.workspace_listbox.selection_set(0)
                    self.current_workspace = self.workspace_listbox.get(0)
                
                # 刷新当前工作区的文档
                self.refresh_documents()
            else:
                messagebox.showerror("Error", f"Failed to fetch workspaces: {response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def delete_workspace(self):
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace to delete")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        if messagebox.askyesno("Confirm", f"Are you sure you want to delete workspace '{workspace_name}'?"):
            api_key = self.api_key_entry.get()
            url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}"
            headers = {"Authorization": f"Bearer {api_key}"}
            
            try:
                response = requests.delete(url, headers=headers)
                if response.status_code == 200:
                    messagebox.showinfo("Success", "Workspace deleted successfully!")
                    self.refresh_workspaces()
                else:
                    messagebox.showerror("Error", f"Failed to delete workspace: {response.text}")
            except Exception as e:
                messagebox.showerror("Error", str(e))

    def select_file(self):
        file_path = filedialog.askopenfilename()
        if file_path:
            self.file_label.config(text=os.path.basename(file_path))
            self.upload_file(file_path)

    def upload_file(self, file_path):
        # 检查是否选择了工作区
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace first")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        api_key = self.api_key_entry.get()
        
        # 上传文件
        upload_url = f"{self.api_base_url}/api/v1/document/upload"
        upload_headers = {"Authorization": f"Bearer {api_key}"}
        
        try:
            # 上传文件
            with open(file_path, 'rb') as f:
                files = {'file': f}
                upload_response = requests.post(upload_url, headers=upload_headers, files=files)
                
            if upload_response.status_code == 200:
                upload_data = upload_response.json()
                doc_location = upload_data.get("documents", [{}])[0].get("location", "")
                
                if doc_location:
                    # 构建更新嵌入向量的请求
                    update_url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/update-embeddings"
                    update_headers = {
                        "Authorization": f"Bearer {api_key}",
                        "Content-Type": "application/json"
                    }
                    update_data = {
                        "adds": [doc_location],
                        "deletes": []
                    }
                    
                    # 发送更新嵌入向量的请求
                    update_response = requests.post(update_url, headers=update_headers, json=update_data)
                    
                    if update_response.status_code == 200:
                        messagebox.showinfo("Success", "File uploaded and embeddings updated successfully!")
                    else:
                        messagebox.showwarning("Warning", f"File uploaded but failed to update embeddings: {update_response.text}")
                        
                self.refresh_documents()
            else:
                messagebox.showerror("Error", f"Failed to upload file: {upload_response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def upload_website(self):
        # 检查是否选择了工作区
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace first")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        website = self.website_entry.get()
        if not website:
            messagebox.showwarning("Warning", "Please enter a website URL")
            return
            
        api_key = self.api_key_entry.get()
        upload_url = f"{self.api_base_url}/api/v1/document/upload-link"
        upload_headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        upload_data = {"link": website}
        
        try:
            # 上传网站
            upload_response = requests.post(upload_url, headers=upload_headers, json=upload_data)
            if upload_response.status_code == 200:
                upload_data = upload_response.json()
                doc_location = upload_data.get("documents", [{}])[0].get("location", "")
                
                if doc_location:
                    # 构建更新嵌入向量的请求
                    update_url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/update-embeddings"
                    update_headers = {
                        "Authorization": f"Bearer {api_key}",
                        "Content-Type": "application/json"
                    }
                    update_data = {
                        "adds": [doc_location],
                        "deletes": []
                    }
                    
                    # 发送更新嵌入向量的请求
                    update_response = requests.post(update_url, headers=update_headers, json=update_data)
                    
                    if update_response.status_code == 200:
                        messagebox.showinfo("Success", "Website uploaded and embeddings updated successfully!")
                    else:
                        messagebox.showwarning("Warning", f"Website uploaded but failed to update embeddings: {update_response.text}")
                
                self.refresh_documents()
                self.website_entry.delete(0, tk.END)
            else:
                messagebox.showerror("Error", f"Failed to upload website: {upload_response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def refresh_documents(self):
        api_key = self.api_key_entry.get()
        url = f"{self.api_base_url}/api/v1/documents"
        headers = {"Authorization": f"Bearer {api_key}"}
        
        try:
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                data = response.json()
                documents = data.get("documents", [])  # 使用 get 方法安全地获取 documents 列表
                self.doc_listbox.delete(0, tk.END)
                for doc in documents:
                    # 使用 title 字段作为文档名称
                    doc_title = doc.get("title", "Unknown Document")
                    self.doc_listbox.insert(tk.END, doc_title)
            else:
                messagebox.showerror("Error", f"Failed to fetch documents: {response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def delete_document(self):
        selection = self.doc_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a document to delete")
            return
            
        doc_title = self.doc_listbox.get(selection[0])
        if messagebox.askyesno("Confirm", f"Are you sure you want to delete document '{doc_title}'?"):
            api_key = self.api_key_entry.get()
            url = f"{self.api_base_url}/api/v1/document/{doc_title}"
            headers = {"Authorization": f"Bearer {api_key}"}
            
            try:
                response = requests.delete(url, headers=headers)
                if response.status_code == 200:
                    messagebox.showinfo("Success", "Document deleted successfully!")
                    self.refresh_documents()
                else:
                    messagebox.showerror("Error", f"Failed to delete document: {response.text}")
            except Exception as e:
                messagebox.showerror("Error", str(e))

    def create_thread(self):
        thread_name = self.thread_entry.get()
        if not thread_name:
            messagebox.showwarning("Warning", "Please enter a thread name")
            return
            
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        api_key = self.api_key_entry.get()
        url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/thread/new"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        data = {
            "name": thread_name,
            "slug": thread_name.lower()
        }
        
        try:
            response = requests.post(url, headers=headers, json=data)
            if response.status_code == 200:
                messagebox.showinfo("Success", "Thread created successfully!")
            else:
                messagebox.showerror("Error", f"Failed to create thread: {response.text}")
        except Exception as e:
            messagebox.showerror("Error", str(e))

    def on_enter(self, event):
        """处理回车键事件"""
        # 检查是否有Shift键按下
        if event.state & 0x1:  # Shift键被按下
            return None  # 让默认行为发生（插入换行）
        self.send_message()
        return "break"  # 阻止默认的回车行为

    def on_shift_enter(self, event):
        """处理Shift+回车键事件"""
        # 在光标位置插入换行符
        self.question_entry.insert(tk.INSERT, "\n")
        return "break"  # 阻止默认的回车行为

    def append_chat(self, text, is_user=False):
        """添加聊天消息到显示区域"""
        self.chat_display.config(state=tk.NORMAL)
        if is_user:
            self.chat_display.insert(tk.END, text + "\n", "user")
        else:
            self.chat_display.insert(tk.END, text + "\n")
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.see(tk.END)

    def update_api_url(self):
        """更新API基础URL"""
        new_url = self.api_url_entry.get().strip()
        if new_url:
            try:
                # 测试新的API URL是否可用
                api_key = self.api_key_entry.get()
                test_url = f"{new_url}/api/v1/workspaces"
                headers = {"Authorization": f"Bearer {api_key}"}
                response = requests.get(test_url, headers=headers)
                
                if response.status_code == 200:
                    self.api_base_url = new_url
                    messagebox.showinfo("Success", "API URL updated successfully!")
                    # 刷新工作区列表
                    self.refresh_workspaces()
                else:
                    messagebox.showerror("Error", f"Failed to connect to new API URL: {response.text}")
                    # 清空工作区列表
                    self.workspace_listbox.delete(0, tk.END)
            except Exception as e:
                messagebox.showerror("Error", f"Failed to connect to new API URL: {str(e)}")
                # 清空工作区列表
                self.workspace_listbox.delete(0, tk.END)
        else:
            messagebox.showwarning("Warning", "Please enter a valid URL")

    def send_message(self):
        question = self.question_entry.get("1.0", tk.END).strip()
        if not question:
            return
            
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        thread_name = self.thread_entry.get()
        if not thread_name:
            messagebox.showwarning("Warning", "Please enter a thread name")
            return
            
        # 检查线程是否存在，如果不存在则创建
        api_key = self.api_key_entry.get()
        url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/thread/{thread_name.lower()}"
        headers = {"Authorization": f"Bearer {api_key}"}
        
        try:
            # 先检查线程是否存在
            response = requests.get(url, headers=headers)
            if response.status_code != 200:
                # 线程不存在，创建新线程
                create_url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/thread/new"
                create_headers = {
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                }
                create_data = {
                    "name": thread_name,
                    "slug": thread_name.lower()
                }
                create_response = requests.post(create_url, headers=create_headers, json=create_data)
                if create_response.status_code != 200:
                    messagebox.showerror("Error", f"Failed to create thread: {create_response.text}")
                    return
            
            # 发送消息
            chat_url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/thread/{thread_name.lower()}/chat"
            chat_headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            chat_data = {
                "message": question,
                "mode": self.mode_var.get()
            }
            
            # 先显示用户消息
            self.append_chat(f"You: {question}", is_user=True)
            self.question_entry.delete("1.0", tk.END)
            
            # 发送请求
            chat_response = requests.post(chat_url, headers=chat_headers, json=chat_data)
            if chat_response.status_code == 200:
                result = chat_response.json()
                answer = result['textResponse']
                if '</think>' in answer:  # 只有在存在 think 标签时才分割
                    answer = answer.split('</think>')[-1].strip()
                sources = result.get('sources', [])
                
                # 显示AI回复
                self.append_chat(f"AI: {answer}")
                
                # 检查是否需要切换模式
                if answer == "There is no relevant information in this workspace to answer your query.":
                    if self.mode_var.get() == "query":
                        if messagebox.askyesno("Mode Switch", "当前查询模式无法获取相关信息。是否切换到聊天模式？"):
                            self.mode_var.set("chat")
                            self.append_chat("\n已切换到聊天模式，请重新发送消息。")
                            # 重新选择之前的工作区
                            for i in range(self.workspace_listbox.size()):
                                if self.workspace_listbox.get(i) == workspace_name:
                                    self.workspace_listbox.selection_clear(0, tk.END)
                                    self.workspace_listbox.selection_set(i)
                                    break
                
                if sources:
                    unique_sources = {src['title']: src for src in sources}.values()
                    source_titles = "\n".join([src['title'] for src in unique_sources])
                    self.append_chat(f"\nSources:\n{source_titles}")
            else:
                self.append_chat(f"Error: {chat_response.text}")
        except Exception as e:
            self.append_chat(f"Error: {str(e)}")

    def load_chat_history(self):
        """加载历史消息"""
        try:
            selection = self.workspace_listbox.curselection()
            if not selection:
                messagebox.showwarning("Warning", "请先选择工作区")
                return
                
            workspace_name = self.workspace_listbox.get(selection[0])
            thread_name = self.thread_entry.get()
            
            api_key = self.api_key_entry.get()
            url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/thread/{thread_name.lower()}/chats"
            headers = {"Authorization": f"Bearer {api_key}"}
            
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                data = response.json()
                history = data.get("history", [])  # 从返回的数据中获取history数组
                self.chat_display.config(state=tk.NORMAL)
                self.chat_display.delete("1.0", tk.END)
                
                for chat in history:
                    if chat.get("role") == "user":
                        # 显示用户消息
                        self.chat_display.insert(tk.END, f"You: {chat['content']}\n", "user")
                    elif chat.get("type") == "chat" and chat.get("role") == "assistant":
                        # 显示AI回复
                        answer = chat['content']
                        if '</think>' in answer:  # 只有在存在 think 标签时才分割
                            answer = answer.split('</think>')[-1].strip()
                        self.chat_display.insert(tk.END, f"AI: {answer}\n")
                        # 显示来源
                        if 'sources' in chat and chat['sources']:
                            sources = chat['sources']
                            unique_sources = {src['title']: src for src in sources}.values()
                            source_titles = "\n".join([src['title'] for src in unique_sources])
                            self.chat_display.insert(tk.END, f"\nSources:\n{source_titles}\n")
                
                self.chat_display.config(state=tk.DISABLED)
                self.chat_display.see(tk.END)
            else:
                messagebox.showerror("Error", f"Failed to load chat history: {response.text}")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to load chat history: {str(e)}")

    def save_chat_history(self):
        """保存对话记录到文件"""
        # 获取当前工作区和线程名称
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "请先选择工作区")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        thread_name = self.thread_entry.get()
        
        # 创建文件名（使用工作区和线程名称）
        default_filename = f"chat_history_{workspace_name}_{thread_name}.txt"
        
        # 让用户选择保存位置
        file_path = filedialog.asksaveasfilename(
            defaultextension=".txt",
            initialfile=default_filename,
            filetypes=[("Text files", "*.txt"), ("All files", "*.*")]
        )
        
        if file_path:
            try:
                # 获取对话内容
                chat_content = self.chat_display.get("1.0", tk.END)
                
                # 添加元数据
                metadata = f"""对话记录
工作区: {workspace_name}
线程: {thread_name}
时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
模式: {self.mode_var.get()}

对话内容:
-------------------
"""
                
                # 写入文件
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(metadata)
                    f.write(chat_content)
                
                messagebox.showinfo("Success", "对话记录已保存！")
            except Exception as e:
                messagebox.showerror("Error", f"保存对话记录失败: {str(e)}")

    def update_embeddings(self):
        """更新选中工作区的文档嵌入向量"""
        # 检查是否选择了工作区
        selection = self.workspace_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a workspace first")
            return
            
        workspace_name = self.workspace_listbox.get(selection[0])
        
        # 获取当前工作区的所有文档
        api_key = self.api_key_entry.get()
        docs_url = f"{self.api_base_url}/api/v1/documents"
        docs_headers = {"Authorization": f"Bearer {api_key}"}
        
        try:
            # 获取所有文档
            docs_response = requests.get(docs_url, headers=docs_headers)
            if docs_response.status_code == 200:
                data = docs_response.json()
                documents = data.get("documents", [])
                
                # 收集所有文档的位置
                doc_locations = []
                for doc in documents:
                    location = doc.get("location")
                    if location:
                        doc_locations.append(location)
                
                if not doc_locations:
                    messagebox.showinfo("Info", "No documents found to update")
                    return
                
                # 发送更新嵌入向量的请求
                update_url = f"{self.api_base_url}/api/v1/workspace/{workspace_name.lower()}/update-embeddings"
                update_headers = {
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json"
                }
                update_data = {
                    "adds": doc_locations,
                    "deletes": []
                }
                
                update_response = requests.post(update_url, headers=update_headers, json=update_data)
                if update_response.status_code == 200:
                    messagebox.showinfo("Success", "Documents embeddings updated successfully!")
                else:
                    messagebox.showerror("Error", f"Failed to update embeddings: {update_response.text}")
            else:
                messagebox.showerror("Error", f"Failed to fetch documents: {docs_response.text}")
        except Exception as e:
            messagebox.showerror("Error", f"Error updating embeddings: {str(e)}")

class WorkspaceConfigDialog(tk.Toplevel):
    def __init__(self, parent):
        super().__init__(parent)
        self.title("Workspace Configuration")
        self.result = None
        
        # 创建主滚动框架
        main_canvas = tk.Canvas(self)
        scrollbar = ttk.Scrollbar(self, orient="vertical", command=main_canvas.yview)
        scrollable_frame = ttk.Frame(main_canvas)

        scrollable_frame.bind(
            "<Configure>",
            lambda e: main_canvas.configure(scrollregion=main_canvas.bbox("all"))
        )

        main_canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        main_canvas.configure(yscrollcommand=scrollbar.set)
        
        # 设置对话框大小和位置
        self.geometry("500x600")
        self.resizable(False, False)
        
        # 创建表单
        tk.Label(scrollable_frame, text="Workspace Name:").pack(anchor="w", padx=10, pady=5)
        self.name_entry = tk.Entry(scrollable_frame)
        self.name_entry.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="Similarity Threshold (0.0-1.0):").pack(anchor="w", padx=10, pady=5)
        self.similarity_entry = tk.Entry(scrollable_frame)
        self.similarity_entry.insert(0, "0.7")
        self.similarity_entry.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="OpenAI Temperature (0.0-1.0):").pack(anchor="w", padx=10, pady=5)
        self.temp_entry = tk.Entry(scrollable_frame)
        self.temp_entry.insert(0, "0.7")
        self.temp_entry.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="OpenAI History Length:").pack(anchor="w", padx=10, pady=5)
        self.history_entry = tk.Entry(scrollable_frame)
        self.history_entry.insert(0, "20")
        self.history_entry.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="OpenAI Prompt:").pack(anchor="w", padx=10, pady=5)
        self.prompt_text = tk.Text(scrollable_frame, height=5)
        self.prompt_text.insert("1.0", "Given the following conversation, relevant context, and a follow up question, reply with an answer to the current question the user is asking. Return only your response to the question given the above information following the users instructions as needed.")
        self.prompt_text.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="Query Refusal Response:").pack(anchor="w", padx=10, pady=5)
        self.refusal_text = tk.Text(scrollable_frame, height=3)
        self.refusal_text.insert("1.0", "There is no relevant information in this workspace to answer your query.")
        self.refusal_text.pack(fill="x", padx=10, pady=5)
        
        tk.Label(scrollable_frame, text="Chat Mode:").pack(anchor="w", padx=10, pady=5)
        self.mode_var = tk.StringVar(value="chat")
        tk.Radiobutton(scrollable_frame, text="Chat", variable=self.mode_var, value="chat").pack(anchor="w", padx=20)
        tk.Radiobutton(scrollable_frame, text="Query", variable=self.mode_var, value="query").pack(anchor="w", padx=20)
        
        tk.Label(scrollable_frame, text="Top N Results:").pack(anchor="w", padx=10, pady=5)
        self.topn_entry = tk.Entry(scrollable_frame)
        self.topn_entry.insert(0, "4")
        self.topn_entry.pack(fill="x", padx=10, pady=5)
        
        # 按钮区域 - 使用固定的底部框架
        button_frame = tk.Frame(self)
        
        # 创建和取消按钮
        tk.Button(button_frame, text="Create", command=self.on_create, width=10).pack(side="right", padx=5)
        tk.Button(button_frame, text="Cancel", command=self.on_cancel, width=10).pack(side="right", padx=5)
        
        # 布局主要组件
        main_canvas.pack(side="left", fill="both", expand=True, padx=5, pady=5)
        scrollbar.pack(side="right", fill="y")
        button_frame.pack(side="bottom", fill="x", padx=10, pady=10)
        
        # 设置模态
        self.transient(parent)
        self.grab_set()
        
        # 绑定鼠标滚轮事件
        self.bind_all("<MouseWheel>", lambda e: main_canvas.yview_scroll(int(-1*(e.delta/120)), "units"))
        
    def validate_inputs(self):
        try:
            name = self.name_entry.get().strip()
            if not name:
                raise ValueError("Workspace name is required")
                
            similarity = float(self.similarity_entry.get())
            if not 0 <= similarity <= 1:
                raise ValueError("Similarity threshold must be between 0 and 1")
                
            temp = float(self.temp_entry.get())
            if not 0 <= temp <= 1:
                raise ValueError("Temperature must be between 0 and 1")
                
            history = int(self.history_entry.get())
            if history < 1:
                raise ValueError("History length must be positive")
                
            topn = int(self.topn_entry.get())
            if topn < 1:
                raise ValueError("Top N must be positive")
                
            return True
        except ValueError as e:
            messagebox.showerror("Validation Error", str(e))
            return False
            
    def on_create(self):
        if self.validate_inputs():
            self.result = {
                "name": self.name_entry.get().strip(),
                "similarityThreshold": float(self.similarity_entry.get()),
                "openAiTemp": float(self.temp_entry.get()),
                "openAiHistory": int(self.history_entry.get()),
                "openAiPrompt": self.prompt_text.get("1.0", "end-1c"),
                "queryRefusalResponse": self.refusal_text.get("1.0", "end-1c"),
                "chatMode": self.mode_var.get(),
                "topN": int(self.topn_entry.get())
            }
            self.destroy()
            
    def on_cancel(self):
        self.result = None
        self.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = ChatApp(root)
    root.mainloop() 